from PIL import Image
from tools.second_det import Second_Detec
from tools.ssd import SSD  # 一次检测
from tools.ssd1 import SSD1  # 二次检测
import matplotlib.pyplot as plt
import cv2
import numpy as np
import datetime
from mindspore import context, Tensor
from mindspore.train.serialization import load


def Detection(opt):
    predict_img = opt.image
    predict_video = opt.video
    if opt.model == 'OM':
        context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
        model_weights = load(opt.target_weights)
        model_helmet = load(opt.target_helmet)
        input_data = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
        input_tensor = Tensor(input_data)
        output_tensor_weights = model_weights(input_tensor)
        output_tensor_helmet = model_helmet(output_tensor_weights)
        output_data = output_tensor_helmet.asnumpy()
    if output_data.name == 'no_helmet':
        # 获取当前时间
        current_time = datetime.datetime.now()
        formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S")
        # 将记录写入data/data.txt文件
        with open('output/data/data.txt', 'a') as file:
            file.write(f"{formatted_time} - 检测到没有佩戴头盔\n")
        print("检测到未戴头盔！警告！")

    if opt.model == 'SSD':
        ssd = SSD(opt)  # 一次检测
        ssd1 = SSD1(opt)  # 二次检测
        # 图像测试
        if predict_img:
            while True:
                img = input('Input image filename(q = quit):')
                if img == 'q':
                    break
                try:
                    image = cv2.imread('input/JPEGImages/' + img + '.jpg')
                except:
                    print('Open Error! Try again!')
                    continue
                else:
                    r_image, axis = ssd.detect_image(image)  # r_image为array,axis[y1,x1,y2,x2]
                    img_reco_crop = Second_Detec(r_image, axis[0], axis[1], axis[2], axis[3])
                    second_res, label = ssd1.detect_image(img_reco_crop)
                    r_image = Image.fromarray(cv2.cvtColor(r_image, cv2.COLOR_BGR2RGB))
                    second_res = Image.fromarray(cv2.cvtColor(second_res, cv2.COLOR_BGR2RGB))
                    if label == 'no_helmet':
                        # 获取当前时间
                        current_time = datetime.datetime.now()
                        formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S")
                        # 将记录写入data/data.txt文件
                        with open('output/data/data.txt', 'a') as file:
                            file.write(f"{formatted_time} - 检测到没有佩戴头盔\n")
                        cv2.imwrite('output/no_helmet/no_helmet_' + img + '.jpg', img_reco_crop)
                        print("检测到未戴头盔！警告！")
                    r_image.save("output/jpg/first_" + img + ".jpg", quality=100)
                    second_res.save("output/jpg/second_" + img + ".jpg", quality=100)
                    plt.subplot(1, 2, 1)
                    plt.imshow(r_image)
                    plt.subplot(1, 2, 2)
                    plt.imshow(second_res, interpolation='none')
                plt.savefig("output/jpg/final_" + img + ".jpg")
                plt.show()
        # 视频测试
        if predict_video:
            current_time = datetime.datetime.now()
            formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S")
            with open('output/data/data.txt', 'w') as file:
                file.write(f"{formatted_time} - 开始检测\n")
            video_path = opt.video_path
            if video_path == '0':
                video_path = int(video_path)
            capture = cv2.VideoCapture(video_path)

            capture.set(cv2.CAP_PROP_BUFFERSIZE, 2)
            capture.set(cv2.CAP_PROP_POS_FRAMES, 15)
            fps = capture.get(cv2.CAP_PROP_FPS)
            # fps = 0.0
            fourcc = cv2.VideoWriter_fourcc(*'XVID')
            size = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))

            out = cv2.VideoWriter('output/video/test.avi', fourcc, fps, size)
            img_reco_crop = None
            Frame_jiange = 14  # 每15帧取一幅图片
            NumFrame = 0
            counter = 0
            while (True):
                # t1 = time.time()
                # 读取某一帧
                res, frame = capture.read()
                if res != True:
                    break
                if res == True:
                    NumFrame = NumFrame + 1
                    if NumFrame % Frame_jiange == 0:
                        result = ssd.detect_image(frame)
                        if type(result) is tuple:
                            frame, axis = result[0], result[1]  # frame为opencv格式，axis[y1,x1,y2,x2]
                            # img_reco_crop = Second_Detec(frame, axis[0], axis[1], axis[2], axis[3])  # 返回np.array 得到截取后的图像
                            img_reco_crop, label = ssd1.detect_image(
                                Second_Detec(frame, axis[0], axis[1], axis[2], axis[3]))
                        if img_reco_crop is not None:
                            frame[(size[1] - axis[2] + axis[0] - 1):(size[1] - 1),
                            (size[0] - axis[3] + axis[1] - 1):(size[0] - 1), :] = img_reco_crop
                            cv2.rectangle(frame, (size[0] - axis[3] + axis[1] - 1, size[1] - axis[2] + axis[0] - 1),
                                          (size[0] - 1, size[1] - 1), (225, 66, 66),
                                          1)
                        if label == 'no_helmet':
                            counter += 1
                            # 获取当前时间
                            current_time = datetime.datetime.now()
                            formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S")
                            # 将记录写入data/data.txt文件
                            with open('output/data/data.txt', 'a') as file:
                                file.write(f"{formatted_time} - 检测到没有佩戴头盔\n")
                            cv2.imwrite('output/no_helmet/no_helmet_' + str(counter) + '.jpg', img_reco_crop)
                            print("检测到未戴头盔！警告！")
                            # 可在此添加告警系统，播放声音等
                        cv2.imshow("video", frame)
                        out.write(frame)
                        if cv2.waitKey(1) & 0xff == ord('q'):
                            capture.release()
                            cv2.destroyAllWindows()
                            break
